from transformers import AutoImageProcessor, AutoModelForImageClassification
import torch
from datasets import load_dataset
import gradio as gr

dataset = load_dataset("huggingface/cats-image")
image = dataset["test"]["image"][0]

image_processor = AutoImageProcessor.from_pretrained("microsoft/resnet-18")
model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-18")

def foward(image):
    inputs = image_processor(image, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits
    idx = logits.argmax(-1).item()
    return model.config.id2label[idx]

# model predicts one of the 1000 ImageNet classes
demo = gr.Interface(fn=foward, inputs="image", outputs="text")
demo.launch()